package cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V1;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.util.Arrays;

/**
 * 初始解析卷积神经网络CNN
 */
public class MnistPredict {
    public static void main(String[] args) throws Exception {
        // 加载预训练好的模型
        MnistPredictor predictor = new MnistPredictor();

        // 加载图片
        BufferedImage img = ImageIO.read(new File(MnistTrain.BASE_PATH + "/mnist_png/decompression/25785.png"));
        int width = img.getWidth();
        int height = img.getHeight();
 
        // 将RGB图片变为灰度图，并将每个像素存于二维数组
        int[][] grayArray = new int[height][width];
        for (int i = 0; i < height; i++) {
            for (int j = 0; j < width; j++) {
                int rgb = img.getRGB(j, i);
                int gray = (int) (0.2989 * ((rgb >> 16) & 0xff) + 0.5870 * ((rgb >> 8) & 0xff) + 0.1140 * (rgb & 0xff));
                // 二值化
                if (gray <= 120)
                    grayArray[i][j] = 255;
            }
        }

        // 标准化和类型转换
//        float[][] normalizedArray = new float[height][width];
//        for (int i = 0; i < height; i++) {
//            for (int j = 0; j < width; j++) {
//                normalizedArray[i][j] = grayArray[i][j] / 255.0f;
//            }
//        }
//
//        // 将像素值转换为uint8类型
//        byte[][] uint8Array = new byte[height][width];
//        for (int i = 0; i < height; i++) {
//            for (int j = 0; j < width; j++) {
//                uint8Array[i][j] = (byte) (normalizedArray[i][j] * 255);
//            }
//        }

        // 对当前的二维数组resize，因为训练好的模型要求输入图像大小为28×28
        int[][] expandedSubArray = resizeImage2(grayArray, 28, 28);
        System.out.println("传递给predict方法的图像数组：" + Arrays.deepToString(expandedSubArray));
 
        // 保存，看看处理后的图像是什么样的，检查有没有处理错误
        save(expandedSubArray, MnistTrain.BASE_PATH + "/imgResult.jpg");

        // 预测
        int prediction = predictor.predict(expandedSubArray);
 
        System.out.println("The input image is predicted to be digit " + prediction);
    }
 
    private static void save(int[][] data, String path) {
        int width = data[0].length;
        int height = data.length;
        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY);
 
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int gray = data[y][x];
                int rgb = (gray << 16) | (gray << 8) | gray;
                image.setRGB(x, y, rgb);
            }
        }
 
        try {
            ImageIO.write(image, "png", new File(path));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    private static int[][] resizeImage2(int[][] input, int newWidth, int newHeight) {
        int height = input.length;
        int width = input[0].length;
        float widthRatio = (float) width / newWidth;
        float heightRatio = (float) height / newHeight;
        int[][] output = new int[newHeight][newWidth];

        for (int y = 0; y < newHeight; y++) {
            for (int x = 0; x < newWidth; x++) {
                float px = x * widthRatio;
                float py = y * heightRatio;
                int x1 = (int) Math.floor(px);
                int y1 = (int) Math.floor(py);
                int x2 = Math.min(x1 + 1, width - 1);
                int y2 = Math.min(y1 + 1, height - 1);

                float dx = px - x1;
                float dy = py - y1;

                float topLeft = input[y1][x1];
                float topRight = input[y1][x2];
                float bottomLeft = input[y2][x1];
                float bottomRight = input[y2][x2];

                float topAvg = topLeft + dx * (topRight - topLeft);
                float bottomAvg = bottomLeft + dx * (bottomRight - bottomLeft);

                output[y][x] = (int) (topAvg + dy * (bottomAvg - topAvg));
            }
        }

        return output;
    }
    private static int[][] resizeImage3(byte[][] input, int newWidth, int newHeight) {
        int height = input.length;
        int width = input[0].length;
        float widthRatio = (float) width / newWidth;
        float heightRatio = (float) height / newHeight;
        int[][] output = new int[newHeight][newWidth];

        for (int y = 0; y < newHeight; y++) {
            for (int x = 0; x < newWidth; x++) {
                float px = x * widthRatio;
                float py = y * heightRatio;
                int x1 = (int) Math.floor(px);
                int y1 = (int) Math.floor(py);
                int x2 = Math.min(x1 + 1, width - 1);
                int y2 = Math.min(y1 + 1, height - 1);

                float dx = px - x1;
                float dy = py - y1;

                byte topLeft = input[y1][x1];
                byte topRight = input[y1][x2];
                byte bottomLeft = input[y2][x1];
                byte bottomRight = input[y2][x2];

                float topAvg = (topLeft & 0xFF) + dx * ((topRight & 0xFF) - (topLeft & 0xFF));
                float bottomAvg = (bottomLeft & 0xFF) + dx * ((bottomRight & 0xFF) - (bottomLeft & 0xFF));

                output[y][x] = (int) (topAvg + dy * (bottomAvg - topAvg));
            }
        }

        return output;
    }

}
 